package com.lujianfei.imagesearch.utils;

import java.awt.image.BufferedImage;
import java.io.BufferedInputStream;
import java.io.FileInputStream;
import java.io.IOException;
import java.util.ArrayList;
import java.util.List;

public class MnistReader {

	public static class MnistInfo {
		public int number;
		public int width;
		public int height;
		public boolean success;
		public String message;
	}
	/**
	 * get images of 'train' or 'test'
	 * 
	 * @param fileNam
	 *            the file of 'train' or 'test' about image
	 * @return one row show a `picture`
	 */
	private int[][] getImages(String fileName,MnistInfo info) {
		int[][] x = null;
		try (BufferedInputStream bin = new BufferedInputStream(new FileInputStream(fileName))) {
			byte[] bytes = new byte[4];
			bin.read(bytes, 0, 4);
			if (2051 != bytesToDecimal(bytes)) { // 读取魔数
				info.success = false;
				info.message = "文件不符合要求";
			} else {
				bin.read(bytes, 0, 4);
				int number = bytesToDecimal(bytes); // 读取样本总数
				info.number = number;
				bin.read(bytes, 0, 4);
				int xPixel = bytesToDecimal(bytes); // 读取每行所含像素点数
				info.width = xPixel;
				bin.read(bytes, 0, 4);
				int yPixel = bytesToDecimal(bytes); // 读取每列所含像素点数
				info.height = yPixel;
				x = new int[number][xPixel * yPixel];
				for (int i = 0; i < number; i++) {
					int[] element = new int[xPixel * yPixel];
					for (int j = 0; j < xPixel * yPixel; j++) {
						element[j] = bin.read(); // 逐一读取像素值
					}
					x[i] = element;
				}
				info.success = true;
			}
			
		} catch (IOException e) {
			info.success = false;
			info.message = "" + e;
		}
		return x;
	}

	/**
	 * get labels of `train` or `test`
	 *
	 * @param fileName
	 *            the file of 'train' or 'test' about label
	 * @return
	 */
	public int[] getLabels(String fileName,MnistInfo info) {
		int[] y = null;
		try (BufferedInputStream bin = new BufferedInputStream(new FileInputStream(fileName))) {
			byte[] bytes = new byte[4];
			bin.read(bytes, 0, 4);
			if (2049 != bytesToDecimal(bytes)) {
				info.success = false;
				info.message = "文件不符合要求";
			} else {
				bin.read(bytes, 0, 4);
				int number = bytesToDecimal(bytes);
				y = new int[number];
				for (int i = 0; i < number; i++) {
					y[i] = bin.read();
				}
			}
			info.success = true;
		} catch (IOException e) {
			info.success = false;
			info.message = "" + e;
		}
		return y;
	}

	public List<BufferedImage> getBufferedImages(String fileName,MnistInfo info) {
		List<BufferedImage> images = new ArrayList<BufferedImage>();
		int[][] buffers = getImages(fileName,info);
		for (int i = 0; i < info.number; i++) {
			int[] b = buffers[i];
			images.add(getImage(b,info.width,info.height));
		}
		return images;
	}

	/**
	 * draw a gray picture and the image format is JPEG.
	 *
	 * @param pixelValues
	 *            pixelValues and ordered by column.
	 * @param width
	 *            width
	 * @param height
	 *            high
	 * @param fileName
	 *            image saved file.
	 */
	private BufferedImage getImage(int[] pixelValues, int width, int height) {
		BufferedImage bufferedImage = new BufferedImage(width, height, BufferedImage.TYPE_INT_RGB);
		for (int i = 0; i < width; i++) {
			for (int j = 0; j < height; j++) {
				int pixel = 0xff - pixelValues[i * height + j];
				int value = pixel + (pixel << 8) + (pixel << 16); // r = g = b 时，正好为灰度
				bufferedImage.setRGB(j, i, value);
			}
		}
		return bufferedImage;
	}

	int bytesToDecimal(byte[] b) {
		int result = 0;
		for (int i = 0; i < b.length; i++) {
			result = result << 8;
			result += b[i];
		}
		return result;
	}
}
